{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class UFS:\n",
    "    def __init__(self, n):\n",
    "        self.n = n\n",
    "        self.f = list(range(n + 1))\n",
    "\n",
    "        self.size = [1] * (n + 1)\n",
    "        self.dist = [0] * (n + 1)\n",
    "\n",
    "    def is_root(self, a):\n",
    "        return a == self.f[a]\n",
    "\n",
    "    def find(self, a):\n",
    "        if not self.is_root(a):\n",
    "            f = self.f[a]\n",
    "\n",
    "            self.f[a] = self.find(self.f[a])\n",
    "            self.dist[a] += self.dist[f]\n",
    "\n",
    "        return self.f[a]\n",
    "\n",
    "    def union(self, a, b, d=0):  # a branch, b root\n",
    "        root_a, root_b = self.find(a), self.find(b)\n",
    "\n",
    "        if root_a != root_b:\n",
    "            self.f[root_a] = root_b\n",
    "\n",
    "            self.dist[root_a] = d\n",
    "            self.size[root_b] += self.size[root_a]\n",
    "\n",
    "    def is_connected(self, a, b):\n",
    "        return self.find(a) == self.find(b)\n",
    "\n",
    "    def get_roots(self):\n",
    "        return [i for i in range(1, self.n + 1) if self.is_root(i)]\n",
    "\n",
    "    def root_dist(self, a):\n",
    "        self.find(a)\n",
    "        return self.dist[a]\n",
    "\n",
    "    def get_dist(self, a, b):\n",
    "        if not self.is_connected(a, b):\n",
    "            return -1\n",
    "\n",
    "        return self.dist[a] - self.dist[b]\n",
    "\n",
    "    def get_size(self, a):\n",
    "        root = self.find(a)\n",
    "        return self.size[root]"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "3a176caed2e623e3a36a87b094ae4f1cc68623d3a88a2d36a2ea84f7c5d7277a"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
